import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from skimage import io, filters, feature
from scipy.ndimage import uniform_filter
import joblib

# 数据准备和预处理
def load_and_preprocess_images(image_path, mask_path):
    image = io.imread(image_path)
    mask = io.imread(mask_path)

    # 确保掩模是单通道
    if mask.ndim == 3:
        mask = mask[:, :, 0]

    # 检查掩模与图像大小一致
    if mask.shape != image.shape[:2]:
        raise ValueError("掩模形状 {} 与图像形状 {} 不匹配".format(mask.shape, image.shape))

    return image, mask  # 返回原始图像和掩模

# 应用不同类型的滤波器
def extract_features(image):
    features = []

    # 均值滤波
    mean_filtered = uniform_filter(image, size=3)
    features.append(mean_filtered.flatten())

    # 高斯滤波
    gaussian_filtered = filters.gaussian(image, sigma=1)
    features.append(gaussian_filtered.flatten())

    # Sobel滤波
    sobel_filtered = filters.sobel(image)
    features.append(sobel_filtered.flatten())

    # Canny边缘检测
    canny_edges = feature.canny(image)
    features.append(canny_edges.flatten())

    return np.array(features).T  # 转置为特征矩阵

# 加载和预处理数据
image_1, mask_1 = load_and_preprocess_images("/shayaprediction\\Sandstone_1.tif",
                                             "/shayaprediction\\Sandstone_1_segment.tif")

# 输出图像形状
assert image_1.shape[:2] == (1024, 996), f"图像形状应为 (1024, 996)，实际为 {image_1.shape}"
print(f"图像形状: {image_1.shape}")

# 特征提取
X = extract_features(image_1)
y = mask_1.flatten()

# 划分训练集和测试集
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
print("完成从砂岩截面图1及其对应分区中获取X和y")
print("完成train_test_split")

# 模型训练
clf = RandomForestClassifier(n_estimators=100, random_state=42)
clf.fit(X_train, y_train)
print("完成随机森林模型clf的训练")

# 评估模型
accuracy = clf.score(X_val, y_val)
print(f"准确率: {accuracy}")

# 保存模型
joblib.dump(clf, 'random_forest_model.joblib')
print("已保存随机森林模型clf到硬盘")